use std::collections::HashMap;
use std::sync::{Arc, Mutex};

#[derive(Debug, Clone, PartialEq)]
pub enum TaskStatus {
    Pending,
    Running,
    Success,
    Failed,
}

#[derive(Clone)]
pub struct TaskManager {
    tasks: Arc<Mutex<HashMap<String, TaskStatus>>>, // 任务 ID 和状态的映射
    total_tasks: Arc<Mutex<usize>>,                 // 任务总数
    completed_tasks: Arc<Mutex<usize>>,
    failed_tasks: Arc<Mutex<usize>>,
}

impl TaskManager {
    pub fn new() -> Self {
        Self {
            tasks: Arc::new(Mutex::new(HashMap::new())),
            total_tasks: Arc::new(Mutex::new(0)),
            completed_tasks: Arc::new(Mutex::new(0)),
            failed_tasks: Arc::new(Mutex::new(0)),
        }
    }

    // 注册新任务
    pub fn register_task(&self, task_id: String) {
        let mut tasks = match self.tasks.lock() {
            Ok(tasks) => tasks,
            Err(_) => panic!("Failed to lock tasks"),
        };
        tasks.insert(task_id, TaskStatus::Pending);

        // 更新任务总数
        let mut total_tasks = match self.total_tasks.lock() {
            Ok(total_tasks) => total_tasks,
            Err(_) => panic!("Failed to lock total_tasks"),
        };
        *total_tasks += 1;
    }

    // 更新任务状态
    pub fn update_task_status(&self, task_id: String, status: TaskStatus) {
        let mut tasks = match self.tasks.lock() {
            Ok(tasks) => tasks,
            Err(_) => panic!("Failed to lock tasks"),
        };
        tasks.insert(task_id, status.clone());

        // 如果任务完成，更新已完成任务数
        if status == TaskStatus::Success || status == TaskStatus::Failed {
            let mut completed_tasks = match self.completed_tasks.lock() {
                Ok(completed_tasks) => completed_tasks,
                Err(_) => panic!("Failed to lock completed_tasks"),
            };
            *completed_tasks += 1;
            if status == TaskStatus::Failed {
                let mut failed_tasks = match self.failed_tasks.lock() {
                    Ok(failed_tasks) => failed_tasks,
                    Err(_) => panic!("Failed to lock failed_tasks"),
                };
                *failed_tasks += 1;
            }
        }
    }

    // 获取任务进度
    pub fn get_progress(&self) -> (usize, usize, usize) {
        let total_tasks = match self.total_tasks.lock() {
            Ok(total_tasks) => total_tasks,
            Err(_) => panic!("Failed to lock total_tasks"),
        };
        let completed_tasks = match self.completed_tasks.lock() {
            Ok(completed_tasks) => completed_tasks,
            Err(_) => panic!("Failed to lock completed_tasks"),
        };
        let failed_tasks = match self.failed_tasks.lock() {
            Ok(failed_tasks) => failed_tasks,
            Err(_) => panic!("Failed to lock failed_tasks"),
        };
        (*total_tasks, *failed_tasks, *completed_tasks)
    }

    // 检查所有任务是否完成
    pub fn all_tasks_completed(&self) -> bool {
        let total_tasks = match self.total_tasks.lock() {
            Ok(total_tasks) => total_tasks,
            Err(_) => panic!("Failed to lock total_tasks"),
        };
        let completed_tasks = match self.completed_tasks.lock() {
            Ok(completed_tasks) => completed_tasks,
            Err(_) => panic!("Failed to lock completed_tasks"),
        };
        *completed_tasks == *total_tasks
    }

    // 停止所有任务
    pub fn stop_all_tasks(&self) {
        let mut tasks = match self.tasks.lock() {
            Ok(tasks) => tasks,
            Err(_) => panic!("Failed to lock tasks"),
        };

        // 遍历所有任务并将其状态更新为 Failed
        for (_, status) in tasks.iter_mut() {
            if *status != TaskStatus::Success && *status != TaskStatus::Failed {
                *status = TaskStatus::Failed;
            }
        }

        // 更新已完成任务数和失败任务数
        let mut completed_tasks = match self.completed_tasks.lock() {
            Ok(completed_tasks) => completed_tasks,
            Err(_) => panic!("Failed to lock completed_tasks"),
        };
        let mut failed_tasks = match self.failed_tasks.lock() {
            Ok(failed_tasks) => failed_tasks,
            Err(_) => panic!("Failed to lock failed_tasks"),
        };

        // 计算需要更新的任务数
        let total_tasks = match self.total_tasks.lock() {
            Ok(total_tasks) => total_tasks,
            Err(_) => panic!("Failed to lock total_tasks"),
        };
        let pending_or_running_tasks = *total_tasks - *completed_tasks;

        // 更新已完成任务数和失败任务数
        *completed_tasks += pending_or_running_tasks;
        *failed_tasks += pending_or_running_tasks;
    }
}
